package com.plant.ai.controller;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;

import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;


@RestController
@RequestMapping("ai")
@CrossOrigin
public class AiController {
    @Resource
    private OllamaChatModel ollamaChatModel;
    /**
     * 流式对话
     *
     * @param message 用户指令
     * @return
     */
    @GetMapping("/streamChat/{message}")
    public Flux<ChatResponse> generateStream(@PathVariable("message") String message) {
        message = "请使用中文简体回答：" + message;
        Prompt prompt = new Prompt(new UserMessage(message));
        return ollamaChatModel.stream(prompt);
    }
    /**
     * 普通对话
     * @param message   用户指令
     * @return
     */
    @GetMapping("/chat/{message}")
    public String generate(@PathVariable("message") String message) {
        message = "请使用中文简体回答：" + message;
        Prompt prompt = new Prompt(new UserMessage(message));
        ChatResponse chatResponse = ollamaChatModel.call(prompt);
        String content = chatResponse.getResult().getOutput().getContent();
        System.out.println("content = " + content);
        return content;
    }
}
